# Copyright 2024 Iguazio
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""add alert activation

Revision ID: 650f0ce2da6f
Revises: 63a7eec6d034
Create Date: 2024-10-30 16:38:07.592754

"""

import os
from datetime import datetime

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import mysql

import mlrun.common.schemas.alert
import mlrun.common.schemas.partition

# revision identifiers, used by Alembic.
revision = "650f0ce2da6f"
down_revision = "63a7eec6d034"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "alert_activations",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column(
            "activation_time", mysql.DATETIME(timezone=True, fsp=3), nullable=False
        ),
        sa.Column(
            "name", sa.String(length=255, collation="utf8mb3_bin"), nullable=False
        ),
        sa.Column(
            "project", sa.String(length=255, collation="utf8mb3_bin"), nullable=False
        ),
        sa.Column("data", sa.JSON(), nullable=True),
        sa.Column(
            "entity_id", sa.String(length=255, collation="utf8mb3_bin"), nullable=False
        ),
        sa.Column(
            "entity_kind",
            sa.String(length=255, collation="utf8mb3_bin"),
            nullable=False,
        ),
        sa.Column(
            "event_kind", sa.String(length=255, collation="utf8mb3_bin"), nullable=False
        ),
        sa.Column(
            "severity", sa.String(length=255, collation="utf8mb3_bin"), nullable=False
        ),
        sa.Column("number_of_events", sa.Integer(), nullable=False),
        sa.PrimaryKeyConstraint("activation_time", "id", name="_alert_activation_uc"),
    )
    op.create_index(
        "ix_alert_activation_activation_time",
        "alert_activations",
        ["activation_time"],
        unique=False,
    )
    op.create_index(
        "ix_alert_activation_project_name",
        "alert_activations",
        ["project", "name"],
        unique=False,
    )
    # ### end Alembic commands ###

    partition_interval = os.getenv("PARTITION_INTERVAL", "YEARWEEK").upper()

    # Validate the partition interval
    if not mlrun.common.schemas.partition.PartitionInterval.is_valid(
        partition_interval
    ):
        raise ValueError(
            f"Partition interval can only be one of the following: "
            f"{mlrun.common.schemas.partition.PartitionInterval.valid_intervals()}"
        )

    # Calculate the date of next partitioning interval
    now_utc = datetime.utcnow()

    partition_interval = mlrun.common.schemas.partition.PartitionInterval(
        partition_interval
    )

    partition_name, partition_value = partition_interval.get_partition_info(now_utc)[0]
    partition_expression = partition_interval.get_partition_expression(
        column_name="activation_time"
    )

    # Construct SQL for partitioning
    partition_sql = f"""
        ALTER TABLE alert_activations
        PARTITION BY RANGE ({partition_expression}) (
            PARTITION {partition_name} VALUES LESS THAN ({partition_value})
        );
    """

    # Execute the SQL command
    op.execute(partition_sql)


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("alert_activations")
    # ### end Alembic commands ###
